from flask import Flask, request, Response, send_from_directory
import torch
import cv2
import numpy as np
import os
import sys

from flask_cors import CORS

# 创建 Flask 应用
app = Flask(__name__, template_folder='.')
# 允许局域网内所有 IP 访问
cors_origins = [
    "http://192.168.*:*",  # 局域网 192.168.x.x 范围
    "http://10.*:*",       # 局域网 10.x.x.x 范围
    "http://localhost:*",  # 本地主机
    "http://127.0.0.1:*"   # 本地回环地址
]

CORS(app, resources={r"/detect": {"origins": cors_origins}})


# 将 yolov5 项目根目录加入 PYTHONPATH
sys.path.append('/Users/mac/PycharmProjects/yolov5-master')

from models.common import DetectMultiBackend
from utils.general import (non_max_suppression, scale_boxes)
from utils.plots import Annotator, colors
from utils.augmentations import letterbox

# 设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载模型
model_path = 'best.pt'
model = DetectMultiBackend(model_path, device=device)
model.eval()

# 图片格式支持
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/')
def index():
    """
    返回 index.html 页面
    """
    return send_from_directory('.', 'index.html')


@app.route('/detect', methods=['POST'])
def detect():
    try:
        if 'file' not in request.files:
            return {"error": "No file provided"}, 400

        file = request.files['file']
        if file.filename == '':
            return {"error": "Empty filename"}, 400

        if not allowed_file(file.filename):
            return {"error": "File type not allowed"}, 400

        # 读取图片
        img_bytes = file.read()
        nparr = np.frombuffer(img_bytes, np.uint8)
        img0 = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        if img0 is None:
            return {"error": "Failed to decode image"}, 400

        # 预处理
        img = letterbox(img0, new_shape=640)[0]
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(device)
        img = img.float() / 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # 推理
        with torch.no_grad():
            pred = model(img)

        # NMS
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)

        # 统计信息
        class_counts = {
            'person': 0,
            'helmet': 0,
            'no-helmet': 0,
            'vest': 0,
            'no-vest': 0,
            'none': 0
        }

        annotator = Annotator(img0, line_width=2, example=str(model.names))
        det = pred[0]
        if len(det):
            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()
            for *xyxy, conf, cls in reversed(det):
                label = f'{model.names[int(cls)]} {conf:.2f}'
                annotator.box_label(xyxy, label, color=colors(int(cls), True))

                cls_name = model.names[int(cls)]
                if cls_name in class_counts:
                    class_counts[cls_name] += 1
                else:
                    print(f"Unknown class: {cls_name}")  # 打印未知类别，便于调试

        # 确保统计信息正确
        print("Class counts:", class_counts)
        result_img = annotator.result()

        # 转换为 JPEG 返回
        _, buffer = cv2.imencode('.jpg', result_img)
        print(f'buffer size: {len(buffer)}')  # 打印 buffer 大小
        if len(buffer) == 0:
            return {"error": "Empty image buffer"}, 500
        headers = {
            'X-Person-Count': str(class_counts['person']),
            'X-Helmet-Count': str(class_counts['helmet']),
            'X-No-Helmet-Count': str(class_counts['no-helmet']),
            'X-Vest-Count': str(class_counts['vest']),
            'X-No-Vest-Count': str(class_counts['no-vest']),
            'X-Total': str(sum(class_counts.values())),
            'Access-Control-Expose-Headers': 'X-Person-Count,X-Helmet-Count,X-No-Helmet-Count,X-Vest-Count,X-No-Vest-Count,X-Total'
        }
        return Response(
            buffer.tobytes(),
            mimetype='image/jpeg',
            headers=headers
        )

    except Exception as e:
        return {"error": str(e)}, 500


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5001, debug=True)
